"""
分布式 ID 生成器模块

提供基于雪花算法的、线程安全的、可注入的分布式 ID 生成服务。
"""
import asyncio
import time
from typing import Dict

from axiom_boot.conf.manager import Settings
from axiom_boot.di import service


@service()
class IdGenerator:
    """
    基于雪花算法的分布式 ID 生成器。

    这是一个可注入的服务，它从应用配置中读取数据中心 ID 和工作节点 ID，
    确保在分布式环境中的唯一性。
    """

    # --- 常量定义 ---
    TIMESTAMP_BITS = 41
    DATACENTER_BITS = 5
    WORKER_BITS = 5
    SEQUENCE_BITS = 12

    MAX_DATACENTER_ID = (1 << DATACENTER_BITS) - 1
    MAX_WORKER_ID = (1 << WORKER_BITS) - 1
    MAX_SEQUENCE = (1 << SEQUENCE_BITS) - 1

    WORKER_ID_SHIFT = SEQUENCE_BITS
    DATACENTER_ID_SHIFT = SEQUENCE_BITS + WORKER_BITS
    TIMESTAMP_LEFT_SHIFT = SEQUENCE_BITS + WORKER_BITS + DATACENTER_BITS
    
    # 起始时间戳 (2024-01-01 00:00:00 UTC)
    EPOCH = 1704067200000

    def __init__(self, settings: Settings):
        """
        初始化 IdGenerator。
        
        Pydantic 在加载配置时已经验证了 datacenter_id 和 worker_id 的范围，
        因此这里不再需要重复检查。
        """
        self.datacenter_id = settings.datacenter_id
        self.worker_id = settings.worker_id
        self.sequence = 0
        self.last_timestamp = -1
        self._lock = asyncio.Lock()

    async def generate_id(self) -> int:
        """异步生成一个唯一的雪花 ID。"""
        async with self._lock:
            timestamp = self._current_timestamp()

            if timestamp < self.last_timestamp:
                raise RuntimeError(f"Clock moved backwards. Refusing to generate id for {self.last_timestamp - timestamp} milliseconds")

            if self.last_timestamp == timestamp:
                self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
                if self.sequence == 0:
                    timestamp = await self._wait_next_millis(self.last_timestamp)
            else:
                self.sequence = 0

            self.last_timestamp = timestamp

            new_id = (
                ((timestamp - self.EPOCH) << self.TIMESTAMP_LEFT_SHIFT) |
                (self.datacenter_id << self.DATACENTER_ID_SHIFT) |
                (self.worker_id << self.WORKER_ID_SHIFT) |
                self.sequence
            )
            return new_id

    def parse_id(self, generated_id: int) -> Dict[str, int]:
        """解析一个雪花 ID，返回其组成部分。"""
        timestamp = ((generated_id >> self.TIMESTAMP_LEFT_SHIFT) + self.EPOCH)
        datacenter_id = (generated_id >> self.DATACENTER_ID_SHIFT) & self.MAX_DATACENTER_ID
        worker_id = (generated_id >> self.WORKER_ID_SHIFT) & self.MAX_WORKER_ID
        sequence = generated_id & self.MAX_SEQUENCE
        return {
            "timestamp_ms": timestamp,
            "datacenter_id": datacenter_id,
            "worker_id": worker_id,
            "sequence": sequence,
        }

    def _current_timestamp(self) -> int:
        """获取当前时间的毫秒时间戳。"""
        return int(time.time() * 1000)

    async def _wait_next_millis(self, last_timestamp: int) -> int:
        """阻塞等待，直到获取到下一毫秒的时间戳。"""
        timestamp = self._current_timestamp()
        while timestamp <= last_timestamp:
            await asyncio.sleep(0.001)  # 等待1毫秒
            timestamp = self._current_timestamp()
        return timestamp 